from dataset import CarvanaDataset
from torch.utils.data import DataLoader
import torch
import torchvision

# num_workers表示线程数，transform表示预处理
def get_loaders(train_dir, train_mask_dir, val_dir, val_mask_dir, batch_size, train_transform, val_transform, num_workers):
    train_set=CarvanaDataset(train_dir, train_mask_dir, train_transform)
    train_loader=DataLoader(train_set,batch_size=batch_size,num_workers=num_workers,shuffle=True)

    val_set=CarvanaDataset(val_dir,val_mask_dir, val_transform)
    val_loader=DataLoader(val_set,batch_size=batch_size,num_workers=num_workers,shuffle=True)
    return train_loader,val_loader

def check_accuracy(loader,model,device):
    num_correct =0
    num_pixels=0
    dice_score=0

    model.eval()    # PyTorch 中用于将模型切换到评估模式的函数
    with torch.no_grad():
        for x,y in loader:
            x=x.to(device)
            y=y.to(device).unsqueeze(1)
            pred=torch.sigmoid(model(x))
            pred=(pred>0.5).float()    #转二值
            num_correct+=(pred==y).sum()
            num_pixels+=torch.numel(pred)   # numel = number pixel，统计y中的像素点个数
            dice_score+=(2*(pred*y).sum())/((pred+y).sum() + 1e-8)

    print(
        f'Got {num_correct}/{num_pixels} with accuracy {num_correct/num_pixels * 100:.2f}%'
    )
    print(f'Dice sorce = {dice_score/len(loader)}')
    model.train()

def save_predictions_as_image(loader,model,device,folder='./saved_val_images/'):
    print('------>Loading predictions')
    model.eval()
    for idx, (x,y) in enumerate(loader):
        x=x.to(device)
        with torch.no_grad():
            pred=torch.sigmoid(model(x))
            pred=(pred>0.5).float()
        torchvision.utils.save_image(pred,f'{folder}/pred_{idx}.png')
        torchvision.utils.save_image(y.unsqueeze(1), f'{folder}/label_{idx}.png')
    model.train()